#pragma once

#include <iostream>
#include <assert.h>
#include <time.h>
using namespace std;

enum Colour // 枚举颜色
{
	RED,
	BLACK
};

template<class K, class V>
struct RBTreeNode
{
	RBTreeNode<K, V>* _left;
	RBTreeNode<K, V>* _right;
	RBTreeNode<K, V>* _parent;

	pair<K, V> _kv;
	Colour _col; // 比AVL树少了平衡因子，多了颜色

	RBTreeNode(const pair<K, V>& kv)
		:_left(nullptr)
		, _right(nullptr)
		, _parent(nullptr)
		, _kv(kv)
	{}
};

template<class K, class V>
struct RBTree
{
	typedef RBTreeNode<K, V> Node;
public:
	bool Insert(const pair<K, V>& kv)
	{
		if (_root == nullptr)
		{
			_root = new Node(kv);
			_root->_col = BLACK; // 根给黑色
			return true;
		}

		Node* parent = nullptr;
		Node* cur = _root;
		while (cur) // 找到要插入的结点
		{
			if (cur->_kv.first < kv.first)
			{
				parent = cur;
				cur = cur->_right;
			}
			else if (cur->_kv.first > kv.first)
			{
				parent = cur;
				cur = cur->_left;
			}
			else
			{
				return false;
			}
		}

		cur = new Node(kv);
		cur->_col = RED; // 默认插入红色结点
		if (parent->_kv.first < kv.first) // 找到位置后插入结点
		{
			parent->_right = cur;
		}
		else
		{
			parent->_left = cur;
		}
		cur->_parent = parent;

		while (parent && parent->_col == RED) // 父亲存在且为红才需要处理
		{
			Node* grandfather = parent->_parent;
			assert(grandfather); // 确定的可以断言下，否则就是插入前就不是红黑树
			assert(grandfather->_col == BLACK);
			if (grandfather->_left == parent)
			{
				Node* uncle = grandfather->_right;

				if (uncle && uncle->_col == RED) // 情况一，叔叔存在且为红(可以直接复制到下面uncle在左边)
				{    // 将父亲和叔叔改为黑，祖父改为红，然后把祖父当成cur，parent变祖父parent继续向上调整。
					parent->_col = uncle->_col = BLACK;
					grandfather->_col = RED;
					cur = grandfather;
					parent = cur->_parent;
				}
				else // 情况二或情况三：叔叔存在且为黑或叔叔不存在
				{
					if(cur == parent->_left) // 情况二的右旋+变色(parent在左)
					{
						//     g      
						//   p   u
						// c
						RotateR(grandfather);
						parent->_col = BLACK; // 父亲变为根了
						grandfather->_col = RED;
					}
					else // 情况二的左右双旋+变色(parent在左)
					{
						//      g      
                        //   p     u
                        //    c
						RotateL(parent);
						RotateR(grandfather);
						cur->_col = BLACK; // cur变为根了
						grandfather->_col = RED;
					}
					break;
				}
			}
			else
			{
				Node* uncle = grandfather->_left;

				if (uncle && uncle->_col == RED) // 情况一，叔叔存在且为红
				{    // 将父亲和叔叔改为黑，祖父改为红，然后把祖父当成cur，parent变祖父parent继续向上调整。
					parent->_col = uncle->_col = BLACK;
					grandfather->_col = RED;
					cur = grandfather;
					parent = cur->_parent;
				}
				else // 情况二或情况三：叔叔存在且为黑或叔叔不存在
				{
					if (cur == parent->_right) // 情况二的左旋+变色(parent在右)
					{
						//     g      
						//   u   p
						//        c
						RotateL(grandfather);
						parent->_col = BLACK; // 父亲变为根了
						grandfather->_col = RED;
					}
					else // 情况二的右左双旋+变色(parent在右)
					{
						//       g      
						//    u     p
						//         c
						RotateR(parent);
						RotateL(grandfather);
						cur->_col = BLACK; // cur变为根了
						grandfather->_col = RED;
					}
					break;
				}
			}
		}
		
		_root->_col = BLACK;
		return true;
	}

	void InOrder()
	{
		_InOrder(_root);
		cout << endl;
	}

	bool IsBalance()
	{
		if (_root == nullptr)
		{
			return true;
		}

		if (_root->_col == RED) // 验证性质二
		{
			cout << "根节点不是黑色" << endl;
			return false;
		}

		int benchmark = 0; // 黑色节点数量基准值
		//Node* cur = _root; // 这种方法是先遍历一遍，然后传值，不过我们可以传引用
		//while (cur)
		//{
		//	if (cur->_col == BLACK)
		//	{
		//		++benchmark;
		//	}
		//	cur = cur->_left;
		//}
		return PrevCheck(_root, 0, benchmark); // 验证性质三和四
	}

protected:
	bool PrevCheck(Node* root, int blackNum, int& benchmark)
	{
		if (root == nullptr)
		{
			if (benchmark == 0)
			{
				benchmark = blackNum;
				return true;
			}

			if (blackNum != benchmark) // 验证性质三
			{
				cout << "某条黑色节点的数量不相等" << endl;
				return false;
			}
			else
			{
				return true;
			}
		}

		if (root->_col == BLACK)
		{
			++blackNum;
		}

		if (root->_col == RED && root->_parent->_col == RED) // 验证性质四
		{
			cout << "存在连续的红色节点" << endl;
			return false;
		}

		return PrevCheck(root->_left, blackNum, benchmark)
			&& PrevCheck(root->_right, blackNum, benchmark);
	}

	void _InOrder(Node* root)
	{
		if (root == nullptr)
		{
			return;
		}
		_InOrder(root->_left);
		cout << root->_kv.first << ":" << root->_kv.second << endl;
		_InOrder(root->_right);
	}

	void RotateL(Node* parent)
	{
		Node* subR = parent->_right; // 动了三个标记了的结点，共更新六个指针，这更新两个指针
		Node* subRL = subR->_left;

		parent->_right = subRL;
		if (subRL) // subRL不为空才更新
		{
			subRL->_parent = parent;
		}

		Node* ppNode = parent->_parent; // 记录parent的parent，防止parent是一颗子树的头结点

		subR->_left = parent; // 再更新两个指针
		parent->_parent = subR;

		if (_root == parent)  // 最后更新两个指针
		{
			_root = subR;
			subR->_parent = nullptr;
		}
		else // parent是一颗子树的头结点
		{
			if (ppNode->_left == parent)
			{
				ppNode->_left = subR;
			}
			else
			{
				ppNode->_right = subR;
			}
			subR->_parent = ppNode;
		}
	}

	void RotateR(Node* parent)
	{
		Node* subL = parent->_left;
		Node* subLR = subL->_right;

		parent->_left = subLR; // 更新两个节点
		if (subLR)
		{
			subLR->_parent = parent;
		}

		Node* ppNode = parent->_parent;

		subL->_right = parent; // 再更新两个节点
		parent->_parent = subL;

		if (_root == parent) // 最后更新两个结点
		{
			_root = subL;
			subL->_parent = nullptr;
		}
		else
		{
			if (ppNode->_left == parent)
			{
				ppNode->_left = subL;
			}
			else
			{
				ppNode->_right = subL;
			}

			subL->_parent = ppNode;
		}
	}

	Node* _root = nullptr;
};